<?php
/**
 * Project: AIGC
 * Author: moufer <moufer@163.com>
 * Date: 2024/7/6 23:41
 */

namespace module\ai\library\channel\llm;

use GuzzleHttp\Client;
use GuzzleHttp\Exception\ClientException;
use GuzzleHttp\Exception\GuzzleException;
use module\ai\library\channel\abstract\Llm;
use module\ai\library\channel\completion\BaiduCompletion;
use module\ai\library\channel\exception\CompletionException;
use module\ai\library\channel\interface\CompletionInterface;
use Psr\Http\Message\ResponseInterface;
use think\facade\Cache;

class BaiduErnieLite extends Llm
{
    protected string $baseUrl = 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat';

    public static function versions(): array
    {
        return [
            '4k'      => 'ERNIE-Lite-4K',
            '8k'      => 'ERNIE-Lite-8K',
            '8k-0922' => 'ERNIE-Lite-8K-0922 (免费)',
            '128k'    => 'ERNIE-Lite-128K',
        ];
    }

    public function getCompletionsUrl(): string
    {
        if (!$endPoint = $this->getEndPoint()) {
            throw new CompletionException('未定义模型地址');
        }

        //拼装地址
        return sprintf(
            '%s/%s?access_token=%s',
            $this->baseUrl, $endPoint, $this->getAccessToken()
        );
    }

    /**
     * 获取请求参数
     * @return array
     */
    protected function getPayload(): array
    {
        $payload = $this->payload;
        $messages = $payload['messages'] ?? null;
        if (empty($messages)) {
            throw new CompletionException('未提供消息内容');
        }
        //取最后一条信息
        $message = $messages[count($messages) - 1];
        if ($message['role'] !== 'user') {
            throw new CompletionException('未提供用户消息');
        }
        //设置系统信息
        if ($this->system) {
            $payload['system'] = $this->system;
        }
        //设置流式回答
        if ($this->isStream) {
            $payload['stream'] = true;
        }
        return $payload;
    }

    /**
     * 获取模型地址
     * @return string|null
     */
    protected function getEndPoint(): ?string
    {
        $model = empty($this->config['model']) ? '8k-0922' : $this->config['model'];
        if (!isset(static::versions()[$model])) {
            throw new CompletionException('未定义百度千帆模型');
        }

        return match ($model) {
            '4k', '128k' => $this->config['model_attr'] ?? '',
            '8k' => 'ernie-lite-8k',
            '8k-0922' => 'eb-instant',
            default => null
        };
    }

    /**
     * @inheritdoc
     */
    protected function setPayload(array $config, bool $override = false): void
    {
        parent::setPayload($config, $override);

        $keys = [
            'temperature',
            'top_p',
            'penalty_score',
            'system',
            'min_output_tokens',
            'max_output_tokens',
            'frequency_penalty',
            'presence_penalty',
            'user_id',
        ];

        foreach ($keys as $key) {
            if (!empty($config[$key])) {
                $this->payload[$key] = $config[$key];
            }
        }

        //payload里不需要的选项
        if (isset($this->payload['model'])) {
            unset($this->payload['model']);
        }
    }

    /**
     * @param string $message
     * @param bool $isStream
     * @param array $clientOptions
     * @inheritdoc
     */
    protected function getResponse(bool $isStream = false, array $clientOptions = []): ResponseInterface
    {
//        try {
        return parent::getResponse($isStream, [
            'headers' => [
                'Content-Type' => 'application/json',
                'Accept'       => 'application/json'
            ]
        ]);
//        } catch (CompletionException $e) {
//            var_dump($e->getCode());
//        }
    }

    /**
     * @inheritdoc
     */
    protected function getCompletionStruct(array $data, bool $isStream = false): ?CompletionInterface
    {
        if (!empty($data['error_code']) && $data['error_code'] === 17) {
            $this->keyModel->save(['status' => -1]);
        }
        return new BaiduCompletion($data, $isStream);
    }

    /**
     * 获取AccessToken
     * @return string
     */
    protected function getAccessToken(): string
    {
        //从缓存里获取去
        $clientId = $this->getKeyModel()?->config['api_key'];
        $clientSecret = $this->getKeyModel()?->config['secret_key'];

        $cacheKey = "ai_access_token:baidu:$clientId";
        $token = Cache::get($cacheKey);
        if (!$token || $token['expire_time']) {
            $url = 'https://aip.baidubce.com/oauth/2.0/token';
            $params = [
                'grant_type'    => 'client_credentials',
                'client_id'     => $clientId,
                'client_secret' => $clientSecret,
            ];
            $client = new Client();
            try {
                $response = $client->post($url, [
                    'form_params' => $params,
                    'headers'     => [
                        'Content-Type' => 'application/json',
                    ]
                ]);
                if ($response->getStatusCode() == 200) {
                    $data = json_decode($response->getBody()->getContents(), true);
                    if (isset($data['error'])) {
                        throw new CompletionException(
                            sprintf('获取AccessToken失败:%s(%s)', $data['error_description'], $data['error'])
                        );
                    }
                    $token = [
                        'access_token' => $data['access_token'],
                        'expire_time'  => time() + $data['expires_in'],
                    ];
                    Cache::set($cacheKey, $token, $data['expires_in']);
                } else {
                    throw new CompletionException(
                        sprintf('获取AccessToken失败:%s', $response->getStatusCode())
                    );
                }
            } catch (ClientException $e) {
                $data = json_decode($e->getResponse()->getBody()->getContents(), true);
                throw new CompletionException($e->getMessage(), $e->getMessage());
            } catch (GuzzleException $e) {
                throw new CompletionException($e->getMessage(), $e->getMessage());
            }
        }

        return $token['access_token'];
    }
}